package cn.daniellee;

import java.io.*;
import java.math.BigDecimal;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.concurrent.CountDownLatch;


public class TcpProxy implements Runnable {

	private static final int SOURCE_PORT = Config.getInt("SourcePort", 25565);

	private static final int TARGET_PORT = Config.getInt("TargetPort", 35565);

	private static final int COUNT_INTERVAL = Config.getInt("CountInerval", 10); //连接次数统计间隔（秒）

	private static final int LIMIT_PER_SECOND = Config.getInt("LimitPerSecond", 3); //开启防御的每秒连接阈值

	private static File listFile = new File("allow.txt");

	private Socket sourceSocket;

	private Socket targetSocket = new Socket();

	private static int secondCount = 0;

	private static boolean enableDefense = false;

	TcpProxy(Socket sourceSocket) {
		this.sourceSocket = sourceSocket;
	}

	private static void log(String message) {
		DateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
		System.out.println(dateFormat.format(new Date()) + " - " + message);
	}

	private static void startServer() {
		try (ServerSocket serverSocket = new ServerSocket(SOURCE_PORT)) {
			Socket socket;
			while ((socket = serverSocket.accept()) != null) {
				secondCount++;
				new Thread(new TcpProxy(socket)).start();
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	@Override
	public void run() {
		BufferedReader reader = null;
		InputStream sourceIn = null, targetIn = null;
		OutputStream sourceOut = null, targetOut = null;
		try {
			boolean allow = !enableDefense;
			if (enableDefense) {
				List<String> ipList = new ArrayList<>();
				reader = new BufferedReader(new InputStreamReader(new FileInputStream(listFile), "UTF-8"));
				String temp;
				while ((temp = reader.readLine()) != null) {
					ipList.add(temp);
				}
				allow = ipList.contains(sourceSocket.getInetAddress().getHostName());
			}
			if (allow) {
				log("Accept: " + sourceSocket.getInetAddress().getHostName());
				sourceIn = sourceSocket.getInputStream();
				sourceOut = sourceSocket.getOutputStream();
				targetSocket.connect(new InetSocketAddress("localhost", TARGET_PORT));
				targetIn = targetSocket.getInputStream();
				targetOut = targetSocket.getOutputStream();
				CountDownLatch countDownLatch = new CountDownLatch(1);
				transfer(sourceIn, targetOut, countDownLatch);
				transfer(targetIn, sourceOut, countDownLatch);
				try {
					countDownLatch.await();
				} catch (InterruptedException e) {
					e.printStackTrace();
				}
			} else {
				log("Denied: " + sourceSocket.getInetAddress().getHostName());
				CountDownLatch countDownLatch = new CountDownLatch(1);
				targetIn = new FileInputStream(new File("motd.txt"));
				sourceOut = sourceSocket.getOutputStream();
				transfer(targetIn, sourceOut, countDownLatch);
				try {
					countDownLatch.await();
				} catch (InterruptedException e) {
					e.printStackTrace();
				}
			}
		} catch (Exception ignored) {
		} finally {
			closeIO(reader);
			closeIO(sourceIn);
			closeIO(sourceOut);
			closeIO(targetIn);
			closeIO(targetOut);
			closeIO(sourceSocket);
		}
	}

	private void transfer(InputStream in, OutputStream out, CountDownLatch latch) {
		new Thread(() -> {
			byte[] bytes = new byte[1024];
			int count;
			try {
				while (0 < (count = in.read(bytes))) {
					out.write(bytes, 0, count);
				}
				out.flush();
			} catch (IOException ignored) {
			} finally {
				if (latch != null) {
					latch.countDown();
				}
			}
		}).start();
	}

	private static void closeIO(Closeable closeable) {
		if (closeable != null) {
			try {
				closeable.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	public static void main(String[] args) {
		if (!listFile.exists()) {
			try {
				listFile.createNewFile();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		java.security.Security.setProperty("networkaddress.cache.ttl", "86400");
		log("转发服务已启动");
		new Thread(() -> {
			try {
				while (true) {
					Thread.sleep(COUNT_INTERVAL * 1000);
					double perSecond = new BigDecimal(secondCount).divide(new BigDecimal(COUNT_INTERVAL), 1, BigDecimal.ROUND_HALF_UP).doubleValue();
					enableDefense = perSecond > LIMIT_PER_SECOND;
					log("当前连接频率：" + perSecond + "次/秒，防御状态：" + (enableDefense ? "开启" : "关闭"));
					secondCount = 0;
				}
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
		}).start();
		TcpProxy.startServer();
	}

}